{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cdfd8725-9a36-419e-9d68-d452baae55c9",
   "metadata": {},
   "source": [
    "# DSPY\n",
    "DSPy: Programming—not prompting—Foundation Models\n",
    "Stanford NLP组的工作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "345598ab-9999-4c78-b921-d577d57765da",
   "metadata": {},
   "source": [
    "## LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "755a89de-52e6-4b63-a01e-c626c45f0086",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cafb8d1e-29a6-497a-a773-aec4de94989b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ollama_mistral = dspy.OllamaLocal(model='mistral-7b-ins')\n",
    "\n",
    "response = ollama_mistral(\"what is digital thread?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "5d48b790-4bec-4cff-83e6-9b337abcf98b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[\"\\n\\nA digital thread refers to the continuous flow of data and information throughout a product's lifecycle, from design and development through manufacturing, supply chain, logistics, maintenance, and disposal. It enables real-time collaboration, transparency, and traceability across different functions and organizations involved in bringing a product to market. By creating a digital thread, companies can improve efficiency, reduce errors, enhance quality, and provide better customer experiences.\\n\\nDigital threads are essential for implementing Industry 4.0 and smart manufacturing initiatives, as they enable the integration of various technologies such as IoT sensors, machine learning algorithms, and augmented reality interfaces to optimize production processes and provide valuable insights into product performance and usage patterns. Additionally\"]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35dd051b-93ed-42ce-b4dd-1f255dc221b2",
   "metadata": {},
   "source": [
    "## Load Env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a4a85f56-257b-49e8-b536-c2bc59e7ead0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\"\"\"\n",
    "Text -> Knowledge Graph\n",
    "1. text -> cypher\n",
    "\n",
    "Constraints:\n",
    "- Use the existing schema before creating new nodes and relationships.\n",
    "\"\"\"\n",
    "import os\n",
    "from dotenv import find_dotenv, load_dotenv\n",
    "\n",
    "# Warning control\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "15ae6f92-fd85-49ae-b0c2-d30f172ccebb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "neo4j://localhost:7687\n"
     ]
    }
   ],
   "source": [
    "load_dotenv(\".env.local\",override=True)\n",
    "NEO4J_URI = os.getenv(\"NEO4J_URI\")\n",
    "NEO4J_USERNAME = os.getenv(\"NEO4J_USERNAME\")\n",
    "NEO4J_PASSWORD = os.getenv(\"NEO4J_PASSWORD\")\n",
    "NEO4J_DATABASE = os.getenv(\"NEO4J_DATABASE\")\n",
    "\n",
    "\n",
    "print(NEO4J_URI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0836828a-e7a5-4e06-bcb5-ee0b378a8ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import neo4j\n",
    "from neo4j import GraphDatabase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "26cbcd49-4292-4789-9f20-b8d37db11862",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_relationships(schema: dict) -> str:\n",
    "    # Parse the JSON string into a Python object if it's not already\n",
    "    if isinstance(schema, str):\n",
    "        data = json.loads(schema)\n",
    "    else:\n",
    "        data = schema\n",
    "    \n",
    "    data = data[0]['relationships']\n",
    "    \n",
    "    # Initialize a list to hold the formatted relationship strings\n",
    "    relationships = []\n",
    "    \n",
    "    # Iterate through each relationship in the data\n",
    "    for relationship in data:\n",
    "        entity1, relation, entity2 = relationship\n",
    "        # Extract the names of the entities and the relationship\n",
    "        entity1_name = entity1['name']\n",
    "        entity2_name = entity2['name']\n",
    "        # Format the string as specified and add it to the list\n",
    "        formatted_relationship = f\"{entity1_name}-{relation}->{entity2_name}\"\n",
    "        relationships.append(formatted_relationship)\n",
    "    \n",
    "    # Join all formatted strings with a newline character\n",
    "    result = \"\\n\".join(relationships)\n",
    "    return result\n",
    "\n",
    "def parse_nodes(schema):\n",
    "    schema = schema\n",
    "    nodes = [node['name'] for node in schema[0]['nodes']]\n",
    "    return \"\\n\".join(nodes)\n",
    "\n",
    "def parse_node_properties(node_properties):\n",
    "    # Initialize a dictionary to accumulate node details\n",
    "    node_details = {}\n",
    "\n",
    "    # Iterate through each item in the input JSON\n",
    "    for item in node_properties:\n",
    "        node_label = item[\"nodeLabels\"][0]  # Assuming there's always one label\n",
    "        prop_name = item[\"propertyName\"]\n",
    "        mandatory = \"required\" if item[\"mandatory\"] else \"optional\"\n",
    "\n",
    "        # Prepare the property string\n",
    "        property_str = f\"{prop_name} ({mandatory})\" if item[\"mandatory\"] else prop_name\n",
    "\n",
    "        # If the node label exists, append the property; otherwise, create a new entry\n",
    "        if node_label in node_details:\n",
    "            node_details[node_label].append(property_str)\n",
    "        else:\n",
    "            node_details[node_label] = [property_str]\n",
    "\n",
    "    # Format the output\n",
    "    output_lines = []\n",
    "    for node, properties in node_details.items():\n",
    "        output_lines.append(f\"{node}\")\n",
    "        for prop in properties:\n",
    "            prop_line = f\"  - {prop}\" if \"required\" in prop else f\"  - {prop}\"\n",
    "            output_lines.append(prop_line)\n",
    "\n",
    "    return \"\\n\".join(output_lines)\n",
    "\n",
    "\n",
    "def parse_rel_properties(rel_properties):\n",
    "    # Initialize a dictionary to accumulate relationship details\n",
    "    rel_details = {}\n",
    "\n",
    "    # Iterate through each item in the input JSON\n",
    "    for item in rel_properties:\n",
    "        # Extract relationship type name, removing :` and `\n",
    "        rel_type = item[\"relType\"][2:].strip(\"`\")\n",
    "        prop_name = item[\"propertyName\"]\n",
    "        mandatory = \"required\" if item[\"mandatory\"] else \"optional\"\n",
    "\n",
    "        # If propertyName is not None, prepare the property string\n",
    "        if prop_name is not None:\n",
    "            property_str = f\"{prop_name} ({mandatory})\"\n",
    "            # If the relationship type exists, append the property; otherwise, create a new entry\n",
    "            if rel_type in rel_details:\n",
    "                rel_details[rel_type].append(property_str)\n",
    "            else:\n",
    "                rel_details[rel_type] = [property_str]\n",
    "        else:\n",
    "            # For relationships without properties, ensure the relationship is listed\n",
    "            rel_details.setdefault(rel_type, [])\n",
    "\n",
    "    # Format the output\n",
    "    output_lines = []\n",
    "    for rel_type, properties in rel_details.items():\n",
    "        output_lines.append(f\"{rel_type}\")\n",
    "        for prop in properties:\n",
    "            output_lines.append(f\"  - {prop}\")\n",
    "\n",
    "    return \"\\n\".join(output_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "beb00926-75f9-4998-8c18-0bcdddf71fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from neo4j import GraphDatabase\n",
    "\n",
    "class Neo4j:\n",
    "    def __init__(self, uri, user: str = None, password: str = None):\n",
    "        self._uri = uri\n",
    "        self._user = user\n",
    "        self._password = password\n",
    "        self._auth = None if (self._user is None and self._password is None) else (self._user, self._password)\n",
    "        self._driver = neo4j.GraphDatabase.driver(self._uri, auth=(self._user, self._password))\n",
    "\n",
    "        # self._driver = neo4j.GraphDatabase.driver\n",
    "\n",
    "        self._verify_connection()\n",
    "\n",
    "    def close(self):\n",
    "        self._driver.close()\n",
    "\n",
    "    def _verify_connection(self):\n",
    "        with self._driver as driver:\n",
    "            driver.verify_connectivity()\n",
    "\n",
    "    def query(self, query, parameters=None, db=None):\n",
    "        assert db is None, \"The Neo4j implementation does not support multiple databases.\"\n",
    "        with self._driver.session(database=db) as session:\n",
    "            result = session.run(query, parameters)\n",
    "            return result.data()\n",
    "\n",
    "    def schema(self, parsed=False):\n",
    "        query = \"\"\"\n",
    "        CALL db.schema.visualization()\n",
    "        \"\"\"\n",
    "        schema = self.query(query)\n",
    "        \n",
    "        if parsed:\n",
    "            return parse_nodes(schema), parse_relationships(schema)\n",
    "        \n",
    "        return schema\n",
    "    \n",
    "    def schema_properties(self, parsed=False):\n",
    "        props = self._schema_node_properties(), self._schema_relationship_properties()\n",
    "        if parsed:\n",
    "            return parse_node_properties(props[0]), parse_rel_properties(props[1])\n",
    "        \n",
    "        return props\n",
    "\n",
    "    def _schema_node_properties(self):\n",
    "        query = \"\"\"\n",
    "        CALL db.schema.nodeTypeProperties()\n",
    "        \"\"\"\n",
    "        return self.query(query)\n",
    "    \n",
    "    def _schema_relationship_properties(self):\n",
    "        query = \"\"\"\n",
    "        CALL db.schema.relTypeProperties()\n",
    "        \"\"\"\n",
    "        return self.query(query)\n",
    "    \n",
    "    def fmt_schema(self):\n",
    "        parsed_schema = self.schema(parsed=True)\n",
    "        parsed_props = self.schema_properties(parsed=True)\n",
    "        parsed = (*parsed_props, parsed_schema[1])\n",
    "        return \"\\n\".join([f\"{element}:\\n{parsed[idx]}\\n\" for idx, element in enumerate([\"NODE LABELS & PROPERTIES\", \"RELATIONSHIP LABELS & PROPERTIES\", \"RELATIONSHIPS\"])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "6984f81f-4305-401a-a279-98d4b20446fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from neo4j_app import Neo4j\n",
    "from neo4j import GraphDatabase\n",
    "# set up Neo4j using NEO4J_URI\n",
    "neo4j = Neo4j(NEO4J_URI,NEO4J_USERNAME,NEO4J_PASSWORD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9b0e8876-e5c5-44e6-a121-9170bac6a3b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dspy.configure(llm=ollama_mistral)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "a74069f4-34c3-4a03-87d3-03a0180ca784",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dspy.settings.lm = ollama_mistral\n",
    "\n",
    "import dsp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "315ce8b5-f922-476a-99f2-77c03a9ff634",
   "metadata": {},
   "outputs": [],
   "source": [
    "dsp.settings.lm = ollama_mistral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "e469bf50-a446-405b-afd2-78688951d3d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'lm': None, 'rm': None, 'branch_idx': 0, 'reranker': None, 'compiled_lm': None, 'force_reuse_cached_compilation': False, 'compiling': False, 'skip_logprobs': False, 'trace': None, 'release': 0, 'log_openai_usage': False, 'bypass_assert': False, 'bypass_suggest': False, 'assert_failures': 0, 'suggest_failures': 0, 'langchain_history': [], 'llm': <dsp.modules.ollama.OllamaLocal object at 0x1685f3790>}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dsp.settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a5c4e637-ec2b-4f54-b1ea-036b883bb1de",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CypherFromText(dspy.Signature):\n",
    "    \"\"\"Instructions:\n",
    "    Create a Cypher MERGE statement to model all entities and relationships found in the text following these guidelines:\n",
    "    - Refer to the provided schema and use existing or similar nodes, properties or relationships before creating new ones.\n",
    "    - Use generic categories for node and relationship labels.\"\"\"\n",
    "\n",
    "    text = dspy.InputField(desc=\"Text to model using nodes, properties and relationships.\")\n",
    "    neo4j_schema = dspy.InputField(desc=\"Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS.\")\n",
    "    statement = dspy.OutputField(desc=\"Cypher statement to merge nodes and relationships found in the text.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "3288fdff-d21e-40a0-8eda-302a7793b4e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_cypher = dspy.ChainOfThought(CypherFromText)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "33caebb8-b686-4414-8ae8-c9b4b9310589",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9622f889-56ec-4b9d-a2bf-1ff33f4933e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "John Singer Sargent (/ˈsɑːrdʒənt/; January 12, 185\n",
      "Lyman Frank Baum (/bɔːm/;[1] May 15, 1856 – May 6,\n",
      "John Davison Rockefeller III (March 21, 1906 – Jul\n",
      "Samuel Langhorne Clemens (November 30, 1835 – Apri\n",
      "Henry Huttleston Rogers (January 29, 1840 – May 19\n",
      "John Pierpont Morgan (April 17, 1837 – March 31, 1\n",
      "John Charles Frémont or Fremont (January 21, 1813 \n",
      "Geronimo (Mescalero-Chiricahua: Goyaałé, Athapasca\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "examples_path = \"./examples/wikipedia-abstracts-v0_0_1.ndjson\"\n",
    "with open(examples_path, \"r\") as f:\n",
    "    # process line by line\n",
    "    for line in f:\n",
    "        data = json.loads(line)\n",
    "        text = data[\"text\"]\n",
    "        print(text[:50])\n",
    "        cypher = generate_cypher(text=text, neo4j_schema=neo4j.fmt_schema())\n",
    "        # neo4j.query(cypher.statement.replace('```', ''))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9397174-57f0-4a72-b7c0-db75b9f486bd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
